package top.lunarye.handler.flowControl.impl;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import top.lunarye.common.domaim.TaskInfo;
import top.lunarye.common.enums.ChannelType;
import top.lunarye.common.enums.RateLimitStrategy;
import top.lunarye.handler.flowControl.FlowControlParam;
import top.lunarye.handler.flowControl.FlowControlService;

/**
 * @author lunar
 * @since 2022-07-20
 */
@Service
@Slf4j
public class FlowControlServiceImpl implements FlowControlService {

    private static final String FLOW_CONTROL_KEY = "flowControl";

    private static final String FLOW_CONTROL_PREFIX = "flow_control_";

    private static final String LIMIT_CONFIG = "{\"flow_control_10\":10, \"flow_control_20\":1}";

    @Override
    public void flowControl(TaskInfo taskInfo, FlowControlParam flowControlParam) {
        RateLimiter rateLimiter = flowControlParam.getRateLimiter();
        Double rateInitValue = flowControlParam.getRateInitValue();

        double costTime = 0;

        // 对比 初始限流值 与 配置限流值，以 配置中心的限流值为准
        Double rateLimitConfig = getRateLimitConfig(taskInfo.getSendChannel());
        if (rateLimitConfig != null && !rateInitValue.equals(rateLimitConfig)) {
            rateLimiter = RateLimiter.create(rateLimitConfig);
            flowControlParam.setRateInitValue(rateLimitConfig);
            flowControlParam.setRateLimiter(rateLimiter);
        }
        // 获取令牌操作,如果不指定令牌数量就默认是一个
        if (RateLimitStrategy.REQUEST_RATE_LIMIT.equals(flowControlParam.getRateLimitStrategy())) {
            costTime = rateLimiter.acquire(1);
        }
        // 指定了需要获取的令牌数量
        if (RateLimitStrategy.SEND_USER_NUM_RATE_LIMIT.equals(flowControlParam.getRateLimitStrategy())) {
            costTime = rateLimiter.acquire(taskInfo.getReceiver().size());
        }

        if (costTime > 0) {
            log.info("consumer {} flow control time {}",
                    ChannelType.getEnumByCode(taskInfo.getSendChannel()).getDescription(), costTime);
        }
    }

    /**
     * 得到限流值的配置
     *
     * @param channelCode
     */
    private Double getRateLimitConfig(Integer channelCode) {
        JSONObject jsonObject = JSON.parseObject(LIMIT_CONFIG);
        if (jsonObject.getDouble(FLOW_CONTROL_PREFIX + channelCode) == null) {
            return null;
        }
        return jsonObject.getDouble(FLOW_CONTROL_PREFIX + channelCode);
    }
}
